#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <pthread.h>
#include "../include/disk_monitor.h"

#define PROC_DISKSTATS "/proc/diskstats"
#define MAX_LINE_LENGTH 256

static pthread_t monitor_thread;
static Disk_IO_Stats current_stats = {0};
static pthread_mutex_t stats_mutex = PTHREAD_MUTEX_INITIALIZER;
static Disk_Monitor_Args* monitor_args = NULL;

// 读取磁盘IO统计信息
static int read_disk_stats(Disk_IO_Stats* stats) {
    FILE* fp;
    char line[MAX_LINE_LENGTH];
    unsigned long long read_sectors, write_sectors;
    unsigned long long read_ops, write_ops;
    
    fp = fopen(PROC_DISKSTATS, "r");
    if (!fp) {
        return -1;
    }
    
    // 重置统计值
    stats->read_bytes = 0;
    stats->write_bytes = 0;
    stats->read_ops = 0;
    stats->write_ops = 0;
    
    // 读取所有磁盘的统计信息
    while (fgets(line, sizeof(line), fp)) {
        if (sscanf(line, "%*d %*d %*s %llu %*u %llu %*u %llu %*u %llu",
                   &read_ops, &read_sectors, &write_ops, &write_sectors) == 4) {
            stats->read_bytes += read_sectors * 512;  // 扇区大小通常为512字节
            stats->write_bytes += write_sectors * 512;
            stats->read_ops += read_ops;
            stats->write_ops += write_ops;
        }
    }
    
    fclose(fp);
    return 0;
}

// 磁盘监控线程函数
void* disk_monitor_thread(void* arg) {
    Disk_Monitor_Args* args = (Disk_Monitor_Args*)arg;
    Disk_IO_Stats prev_stats = {0};
    Disk_IO_Stats curr_stats = {0};
    
    // 首次读取统计信息
    if (read_disk_stats(&prev_stats) < 0) {
        return NULL;
    }
    
    while (args->running) {
        sleep(args->interval);
        
        // 读取当前统计信息
        if (read_disk_stats(&curr_stats) < 0) {
            continue;
        }
        
        // 计算速度
        double interval = args->interval;
        curr_stats.read_speed = (curr_stats.read_bytes - prev_stats.read_bytes) / (1024.0 * interval);
        curr_stats.write_speed = (curr_stats.write_bytes - prev_stats.write_bytes) / (1024.0 * interval);
        
        // 更新统计信息
        pthread_mutex_lock(&stats_mutex);
        current_stats = curr_stats;
        pthread_mutex_unlock(&stats_mutex);
        
        // 保存当前值作为下一次的前值
        prev_stats = curr_stats;
    }
    
    return NULL;
}

// 初始化磁盘监控
int init_disk_monitor(void) {
    monitor_args = malloc(sizeof(Disk_Monitor_Args));
    if (!monitor_args) {
        return -1;
    }
    
    monitor_args->interval = 1;  // 默认1秒更新一次
    monitor_args->running = 1;
    
    if (pthread_create(&monitor_thread, NULL, disk_monitor_thread, monitor_args) != 0) {
        free(monitor_args);
        monitor_args = NULL;
        return -1;
    }
    
    return 0;
}

// 清理磁盘监控
void cleanup_disk_monitor(void) {
    if (monitor_args) {
        monitor_args->running = 0;
        pthread_join(monitor_thread, NULL);
        free(monitor_args);
        monitor_args = NULL;
    }
    
    pthread_mutex_destroy(&stats_mutex);
}

// 获取当前磁盘IO统计信息
Disk_IO_Stats get_disk_io_stats(void) {
    Disk_IO_Stats stats;
    
    pthread_mutex_lock(&stats_mutex);
    stats = current_stats;
    pthread_mutex_unlock(&stats_mutex);
    
    return stats;
} 